#pragma once
#include <vector>
#include "common\define.h"
#include "NvInfer.h"
#include "opencv2/opencv.hpp"
#include "cuda_img_resize.cuh"

namespace key_point
{
	struct DetectRes {
		int classes;
		float x;
		float y;
		float w;
		float h;
		float prob;
		float kps_x;
		float kps_y;
		float kps_s;
	};

class CKeyPoint
{

public:
    CKeyPoint();
    ~CKeyPoint();

	bool Init(int gpu_id, const char* gdats_path);
	void Detect(st_dlm_data* p_imgs, int num, std::vector<std::vector<DetectRes>> &vv_results);
	
private:
	void Reset();
	bool LoadEngine(std::vector<char> v_engine_data);
	std::vector<float> PrepareImage(std::vector<cv::Mat> vec_Mat);
	
	std::vector<std::vector<DetectRes>> PostProcess(const std::vector<cv::Mat>& vec_Mat, st_dlm_data* p_imgs, float* output, const int& outSize);
	void NmsDetect(std::vector<DetectRes>& detections, float nms_threshold);
	float IOUCalculate(const DetectRes& det_a, const DetectRes& det_b);

	st_encrypt_info m_info;
	nvinfer1::ICudaEngine* m_engine = nullptr;
	nvinfer1::IExecutionContext* m_context = nullptr;
	cudaStream_t m_stream; 
	
	std::vector<void*> m_v_dev_buffer;
	std::vector<int64_t> m_v_dev_buf_size;

	std::vector<int> m_strides; 
	int m_num_anchors;

	std::vector<cv::Scalar> m_class_colors;

    std::vector<float> m_img_mean;
    std::vector<float> m_img_std;
    bool dynamic;
  
	int m_out_size[2] = { 0 };
	
	std::vector<float> m_v_out_buf[2];
	std::vector<st_cuda_resize_dev_space> _v_resize_space;

   
};

}//namespace key_point
